# 导入所需库
import pandas as pd
import numpy as np
import statsmodels.api as sm
import matplotlib.pyplot as plt
from sqlalchemy import create_engine

# 数据库连接配置
db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': '123456',
    'database': 'tushare1',
    'port': 3306,
    'charset': 'utf8mb4'
}

# 创建SQLAlchemy引擎
engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}"
    f"@{db_config['host']}:{db_config['port']}/{db_config['database']}"
    f"?charset={db_config['charset']}"
)

# =========== 数据获取部分 ==============
# 定义分块大小
chunk_size = 10000

# 执行SQL查询（修正后的查询语句）
query = """
SELECT d.closes, d.vol, d.amount, 
       m.buy_sm_vol, m.sell_sm_vol, 
       m.buy_md_vol, m.sell_md_vol,
       m.buy_lg_vol, m.sell_lg_vol,
       m.buy_elg_vol, m.sell_elg_vol
FROM date_1 d
JOIN moneyflows m
ON d.ts_code = m.ts_code AND d.trade_date = m.trade_date
WHERE d.trade_date BETWEEN '2023-02-01' AND '2023-07-01' 
  AND d.ts_code='002229.SZ'
"""

# 分块读取数据
chunks = pd.read_sql_query(query, engine, chunksize=chunk_size)
df = pd.concat(chunks, ignore_index=True)

# ==================== 数据预处理部分 ====================
# 计算涨跌幅（修正列名）
df['zd_close'] = round(((df['closes'] - df['closes'].shift(1)) / df['closes'].shift(1)), 2)

# 处理缺失值
df = df.dropna(subset=['zd_close']).reset_index(drop=True)

# ==================== 主成分分析部分 ====================
# 选择数值型自变量（修正列名）
numeric_cols = ['vol', 'amount', 
               'buy_sm_vol', 'sell_sm_vol',
               'buy_md_vol', 'sell_md_vol',
               'buy_lg_vol', 'sell_lg_vol',
               'buy_elg_vol', 'sell_elg_vol']

X = df[numeric_cols].copy()

# 主成分分析
cov_matrix = np.cov(X, rowvar=False)
eigenvalues, eigenvectors = np.linalg.eig(cov_matrix)

# 输出累计贡献率
cumulative_contribution = eigenvalues[:5].sum()/eigenvalues.sum()
print(f"累计贡献率为: {cumulative_contribution:.2%}")

# 选择主成分
n_components = 5
top_eigenvectors = eigenvectors[:, :n_components]

# 计算主成分得分
principal_components = np.dot(X, top_eigenvectors)

# 添加主成分到数据集
pc_cols = [f'PC{i+1}' for i in range(n_components)]
df_pca = pd.concat([df, pd.DataFrame(principal_components, columns=pc_cols)], axis=1)

# ==================== 回归分析部分 ====================
# 准备数据
X_pca = df_pca[['PC1', 'PC3', 'PC4', 'PC5']]  # 根据分析选择主成分
y = df_pca['zd_close']

# 添加常数项
X_pca = sm.add_constant(X_pca)

# 构建模型
model = sm.OLS(y, X_pca)
results = model.fit()

# 输出结果
print("\n回归模型结果:")
print(results.summary())
#============= 可视化部分 ============
# 绘制散点图矩阵
fig, axes = plt.subplots(1, 4, figsize=(18, 5))
for i, col in enumerate(['PC1', 'PC3', 'PC4', 'PC5']):
    axes[i].scatter(X_pca[col], y, s=58, alpha=0.7)
    axes[i].set_xlabel(col)
    axes[i].set_ylabel('涨跌幅')
    axes[i].set_title(f'{col} vs 涨跌幅')
plt.tight_layout()
plt.show()

# ==================== 主成分解释 ============
# 输出主成分方程
for k in range(n_components):
    component_str = f'PC{k+1} = '
    weights = eigenvectors[:, k]
    
    for j, weight in enumerate(weights):
        sign = '+' if weight >= 0 else '-'
        component_str += f' {sign} {abs(weight):.2f}*{numeric_cols[j]}'
    
    print(component_str + '\n')

# 关闭数据库连接
engine.dispose()